#lang racket

(require "util/utils.scm")
(require "util/match-assoc.scm")
(provide opt-known-calls)

(define (analyze-known-functions prg)
  (if (null? prg)
      '()
      (let ((def (car prg)))
	(match def
	  (`((label ,nm) () () (closure () (label ,nm-2)))
	   (cons (cons nm nm-2)
		 (analyze-known-functions (cdr prg))))
	  (else (analyze-known-functions (cdr prg)))))))

(define (opt-known-calls prg)
  (match-assoc prg (definitions code)
    (let ((known (analyze-known-functions code)))
      `((definitions . ,definitions)
	(code . ,(mapply opt-known-calls-procedures code known))))))

(define (opt-known-calls-procedures procedure known-functions)
  (match procedure
    (`((label . ,nm) ,env-1 ,env-2 ,body)
     `((label . ,nm) ,env-1 ,env-2 ,(opt-known-calls-exp body known-functions)))))

(define (opt-known-calls-binding binding known-functions)
  `(,(car binding)
    ,(opt-known-calls-exp (cadr binding) known-functions)))

(define (opt-known-calls-exp exp known-functions)
  (match exp
    (`(quote ,x) `(quote ,x))
    (`(string ,s) `(string ,s))
    (`(atomic ,x) `(atomic ,x))
    (`(builtin ,f . ,args)
     `(builtin ,f . ,(mapply opt-known-calls-exp args known-functions)))
    (`(var ,ty ,x) `(var ,ty ,x))
    ;; the actual optimization here:
    (`(app (var glo ,i) . ,exps)
     (append (cond ((assoc i known-functions) =>
		    (lambda (entry)
		      (list 'known-app (cdr entry))))
		   (else (list 'app `(var glo ,i))))
	     (mapply opt-known-calls-exp exps known-functions)))
    (`(app . ,exps)
     `(app . ,(mapply opt-known-calls-exp exps known-functions)))
    ;;
    (`(begin . ,exps)
     `(begin . ,(mapply opt-known-calls-exp exps known-functions)))
    (`(closure ,env ,lbl)
     `(closure ,env ,lbl))
    (`(if ,t ,c ,a)
     `(if ,(opt-known-calls-exp t known-functions)
	  ,(opt-known-calls-exp c known-functions)
	  ,(opt-known-calls-exp a known-functions)))
    (`(let* ,bindings ,body)
     `(let* ,(mapply opt-known-calls-binding bindings known-functions)
	,(opt-known-calls-exp body known-functions)))
    (else
     (error 'opt-known-calls-exp "no matching pattern" exp))))
